"""测试数据库构建器"""

import numpy as np
import cv2
import tqdm
import argparse
from utils.ellipse import from_ellipse_args, ellipse_fit, ellipse_args, center_ellipse
from database.plain.binad import (
    XuBuilder,
    ChenBuilder,
    ForsythBuilder,
    ChristianBinadBuilder,
)
from database.plain.triad import HanakBuilder, ParkBuilder, ChristianBuilder
from database.hash import XuHashBuilder, ChristianHashBuilder
import os
import time


def arg_parse():
    arg_parser = argparse.ArgumentParser()
    arg_parser.add_argument(
        "--num", type=int, default=30, help="The maximal number of points"
    )
    arg_parser.add_argument("-f", type=float, default=25)
    arg_parser.add_argument("-d", type=float, default=5.5)
    arg_parser.add_argument("--reso", type=int, default=1024)
    arg_parser.add_argument("--noise", type=float, default=0)
    arg_parser.add_argument("--drift", type=float, default=0)
    arg_parser.add_argument("--height", type=float, default=4)
    arg_parser.add_argument("--rotate", type=float, default=np.pi / 5)
    arg_parser.add_argument(
        "--exp", type=int, default=100, help="The number of experiments"
    )
    arg_parser.add_argument("--recall", type=float, default=90)
    arg_parser.add_argument("--outlier", type=float, default=0)
    arg_parser.add_argument(
        "--use_init", action="store_true", help="Use pnp pose as initial"
    )
    arg_parser.add_argument("--debug", action="store_true", help="Debug mode")
    arg_parser.add_argument("--catalog", type=str, default="craters.csv")
    arg_parser.add_argument("--output_dir", type=str, default="./output")
    return arg_parser.parse_args()


if __name__ == "__main__":
    args = arg_parse()
    methodClass = XuBuilder
    name = methodClass.__name__
    builder = methodClass(
        args.catalog, max_field=1000, cache_path=f"./__pycache__/{name}.npz"
    )
    catalog_dict = builder()
    catalog = catalog_dict["catalog"]
    descriptor = catalog_dict["descriptor"]
    idx = catalog_dict["idx"]
    temp = catalog_dict["temp"]
    constants = catalog_dict["constants"]
    builder.to_csv(f"./__pycache__/{name}")
    # 生成虚拟相机矩阵
    K = np.array(
        [
            [args.f / args.d * 1e3, 0, args.reso / 2],
            [0, args.f / args.d * 1e3, args.reso / 2],
            [0, 0, 1],
        ]
    )
    # 平面上任取三个点，构成生成子空间
    M = np.block([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1]])
    # 此处不再是随意选取陨石坑圆，而是从库中读取
    # elli_3d = []
    # C = []
    # cnt_3d = []
    # 生成椭圆的二维点
    x, y, r = temp.T
    points_2d = from_ellipse_args(x, y, r, r, np.zeros_like(x))
    points_2d = np.pad(
        points_2d, ((0, 0), (0, 1), (0, 0)), "constant", constant_values=1
    )
    # 二维点转三维点，以模拟相机成像过程
    elli_3d = np.einsum("ijk,jl->ilk", points_2d, M)
    C = np.array(
        [
            [np.ones_like(x), np.zeros_like(x), -x],
            [np.zeros_like(x), np.ones_like(x), -y],
            [-x, -y, x * x + y * y - r * r],
        ]
    ).T
    cnt_3d = np.array([x, y, np.zeros_like(x)]).T
    # for x, y, r in temp:
    #     points_2d = from_ellipse_args(
    #         x[None], y[None], r[None], r[None], np.zeros_like(x[None])
    #     )
    #     points_2d = np.pad(
    #         points_2d, ((0, 0), (0, 1), (0, 0)), "constant", constant_values=1
    #     )
    #     points_3d = np.einsum("ijk,jl->ilk", points_2d, M)
    #     C.append(np.array([[1, 0, -x], [0, 1, -y], [-x, -y, x * x + y * y - r * r]]))
    #     elli_3d.append(points_3d)
    #     cnt_3d.append((x, y, 0))
    # C = np.array(C)
    # cnt_3d = np.array(cnt_3d)
    # elli_3d = np.concatenate(elli_3d, axis=0)
    output_dir = os.path.join(args.output_dir, f"{args.num}")
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    total = 0
    gbar = tqdm.tqdm(total=args.exp, desc="Testing Catalog")
    while total < args.exp:
        # 生成像平面上的椭圆
        ## 相机随机位姿
        # quat = list(float(x) for x in line.strip().split(",")[4:])
        # R = Rotation.from_quat(quat).as_matrix()
        # Rvec = cv2.Rodrigues(R)[0]
        # position = list(float(x) for x in line.strip().split(",")[1:4])
        Rvec = np.random.randn(3, 1) * args.rotate
        R = cv2.Rodrigues(Rvec)[0]
        position = temp.mean(axis=0) + np.random.randn(3) * temp.std(axis=0)
        position[2] += np.abs(np.random.randn()) * args.height * 1e3
        T = -R @ position
        # P = K @ np.hstack((R, T[:, None]))
        # ## 生成像平面上的椭圆
        # H = P @ M.T
        # C_real = np.linalg.inv(H).T @ C @ np.linalg.inv(H)
        ## 在原始平面上采点，投影至像平面，再拟合为椭圆
        Qs = []
        Us = []
        Qs_prime_idx = []
        obse_Qs_points = []
        points_2d = cv2.projectPoints(cnt_3d.T, Rvec, T, K.astype(np.float32), None)[
            0
        ].squeeze()
        ind = (
            (0 <= points_2d[:, 0])
            & (points_2d[:, 0] < args.reso)
            & (0 <= points_2d[:, 1])
            & (points_2d[:, 1] < args.reso)
        )
        for it, points_3d in enumerate(elli_3d):
            if not ind[it]:
                continue
            points_2d = cv2.projectPoints(
                points_3d[:3].T, Rvec, T, K.astype(np.float32), None
            )[0].squeeze()
            # 加入噪声和漂移
            points_2d += np.random.randn(*points_2d.shape) * args.noise + args.drift
            # 加入异常值点，即误检测的椭圆，同时删除一个正确的椭圆
            if np.random.rand() < args.outlier:
                r = np.random.uniform(10, 100, 2)
                theta = np.random.uniform(0, 2 * np.pi, 1)
                E = from_ellipse_args(
                    *np.block([np.random.uniform(0, 1024, 2), r, theta])
                )
                E, U = ellipse_fit(E[0], E[1])
            else:
                E, U = ellipse_fit(points_2d[:, 0][None], points_2d[:, 1][None])
                # C_real = cone_plane_cross(P.T @ E @ P, (0, 0, 1, 0))
            # 用当前椭圆的参数排除在视场以外的椭圆
            x, y, a, b, theta = ellipse_args(E)
            if x:
                if 0 <= x < args.reso and 0 <= y < args.reso:
                    Qs.append((x, y, a, b, theta))
                    Us.append(U)
                    Qs_prime_idx.append(it)
                    obse_Qs_points.append(points_2d)
        # 考虑真实情况，同一幅图像中椭圆数量少于4个一定测不出位姿，大于20个不符合真实导航场景
        if len(Qs) != args.num:
            continue
        # key = len(Qs)
        Qs = np.concatenate(Qs, axis=1).T
        Us = np.concatenate(Us, axis=0)
        total += 1
        gbar.update(1)
        # 测试图像平面上的椭圆不变量
        with open(f"{output_dir}/arguments.csv", "a") as f:
            print(
                f"{args.f},{args.d},{args.reso}",
                ",".join(map(str, Rvec.squeeze())),
                ",".join(map(str, T.squeeze())),
                file=f,
                sep=",",
            )
        # 再将全部陨石坑写入至文件，作为测试数据
        np.savez(f"{output_dir}/{total}.npz", idx=Qs_prime_idx, points=obse_Qs_points)
        # for it, (Q, C_idx) in enumerate(zip(Qs, Qs_prime_idx)):
        #     invariants = builder.invariant(Q[None], Qs[it + 1 :])
        #     C_ind = Qs_prime_idx
        #     real_invariant = builder.invariant(
        #         C[C_idx, None], C[Qs_prime_idx[it + 1 :], :]
        #     )
        #     ind = np.block(
        #         list(
        #             np.argwhere((idx[0] == C_idx) & (idx[1] == i)).squeeze()
        #             for i in Qs_prime_idx[it + 1 :]
        #         )
        #     )
        #     real_invariant = descriptor[:, ind]
        # 测试三元组型陨石坑不变量
        # for it, (Q, C_idx) in enumerate(zip(Qs, Qs_prime_idx)):
        #     for i in range(it + 1, len(Qs)):
        #         invariants = builder.invariant(Q[None], Qs[i], Qs[i + 1 :])
        #         C_ind = Qs_prime_idx
        #         real_invariant = builder.invariant(
        #             C[C_idx, None],
        #             C[Qs_prime_idx[i], None],
        #             C[Qs_prime_idx[i + 1 :], :],
        #         )
        #         ind = np.block(
        #             list(
        #                 np.argwhere(
        #                     (idx[0] == C_idx)
        #                     & (idx[1] == Qs_prime_idx[i])
        #                     & (idx[2] == ii)
        #                 ).squeeze()
        #                 for ii in Qs_prime_idx[i + 1 :]
        #             )
        #         )
        #         real_invariant = descriptor[:, ind]
